
#include "reg_bw.h"
const size_t MAX_BLOCKS = 8192;
const int32_t THR_SIZE = 64;
__global__ void cuda_add_reg_ori(float *out, float *x, float y, size_t N) {
  size_t num_in_grid = size_t(blockDim.x * gridDim.x * THR_SIZE);
  float tmp[THR_SIZE];
  int tmp2[THR_SIZE];
  for (size_t n = 0; n < N; n += num_in_grid) {
    for (size_t i = 0; i < THR_SIZE; i += 1) {
      size_t idx =
          n + blockIdx.x * blockDim.x * THR_SIZE + threadIdx.x + blockDim.x * i;
      tmp[i] = x[idx];
      tmp2[i] = int(x[idx]);
    }
    int y2 = int(y);
    for (size_t j = 0; j < 1000; j++) {
#pragma unroll
      for (size_t i = 0; i < THR_SIZE; i += 1) {
        tmp2[i] = tmp2[i] + y2;
        tmp[i] = tmp[i] + y;
      }
    }
    for (size_t i = 0; i < THR_SIZE; i += 1) {
      size_t idx =
          n + blockIdx.x * blockDim.x * THR_SIZE + threadIdx.x + blockDim.x * i;
      out[idx] = tmp[i] + tmp2[i];
    }
  }
}

void reg_bw_ori(float *out, float *x, float y, cudaStream_t stream, size_t N) {
  constexpr size_t threads_per_block = 128;
  auto num_blocks = min(size_t((N + threads_per_block * THR_SIZE - 1) /
                               (threads_per_block * THR_SIZE)),
                        MAX_BLOCKS);
  cuda_add_reg_ori<<<num_blocks, threads_per_block, 0, stream>>>(out, x, y, N);
}
